//
// Created by QU on 24-9-21.
//
#include <boost/asio.hpp>
#include <boost/endian/conversion.hpp>

#include <iostream>
#include <vector>
#include <thread>
#include <atomic>
#include <random>
#include <chrono>
#include <future>
#include <iomanip>
#include <MsgNode/Node.h>

class StressClient
{
public:
    struct Config
    {
        int thread_count = 10; // 并发运行的线程数量
        int messages_per_thread = 10000; // 每个线程发送的消息数量
        int small_message_min = 10; // 小消息的最小大小（字节）
        int small_message_max = 1024; // 小消息的最大大小（字节）
        int large_message_min = 1024; // 大消息的最小大小（字节）
        int large_message_max = Node::MAX_PAYLOAD_SIZE; // 大消息的最大大小（字节）
        int multiple_message_count = 10; // 多消息模式下发送的消息数量
        int multiple_message_min = 5; // 多消息模式下每条消息的最小大小（字节）
        int multiple_message_max = 8192; // 多消息模式下每条消息的最大大小（字节）
        int split_message_min = 5; // 分割消息的最小总大小（字节）
        int split_message_max = 8192; // 分割消息的最大总大小（字节）
        int delay_between_messages_min = 0; // 消息间最小延迟时间（毫秒）
        int delay_between_messages_max = 5; // 消息间最大延迟时间（毫秒）
        std::chrono::seconds total_test_duration { 300 }; // 测试的总持续时间（秒）
    };

    StressClient(const std::string& host, const std::string& port, const Config& config = Config { })
        : m_host(host), m_port(port), m_config(config)
    {
    }

    void StartLoad()
    {
        std::vector<std::future<void>> futures;
        std::atomic<bool> should_stop { false };
        auto start_time = std::chrono::steady_clock::now();

        for (int i = 0; i < m_config.thread_count; ++i)
        {
            futures.emplace_back(std::async(std::launch::async, [this, &should_stop, start_time]()
            {
                RunClientThread(should_stop, start_time);
            }));
        }

        for (auto& future : futures)
        {
            future.wait();
        }

        std::cout << "All threads completed. Total threads: " << m_config.thread_count << std::endl;
    }

private:
    void RunClientThread(std::atomic<bool>& should_stop, std::chrono::steady_clock::time_point start_time)
    {
        try
        {
            boost::asio::io_context io_context;
            boost::asio::ip::tcp::socket socket(io_context);
            Connect(socket);

            std::random_device rd;
            std::default_random_engine gen(rd());
            std::uniform_int_distribution<> dis(0, 3);

            for (int i = 0; i < m_config.messages_per_thread && !should_stop; ++i)
            {
                SendRandomMessage(socket, dis(gen));
                DelayBetweenMessages(gen);

                if (std::chrono::steady_clock::now() - start_time >= m_config.total_test_duration)
                {
                    should_stop = true;
                }
            }
        }
        catch (const std::exception& e)
        {
            std::cerr << "Thread " << std::this_thread::get_id() << " exception: " << e.what() << std::endl;
        }
    }

    void Connect(boost::asio::ip::tcp::socket& socket)
    {
        boost::asio::ip::tcp::resolver resolver(socket.get_executor());
        auto endpoints = resolver.resolve(m_host, m_port);
        boost::asio::connect(socket, endpoints);
    }

    void SendRandomMessage(boost::asio::ip::tcp::socket& socket, int message_type)
    {
        switch (message_type)
        {
        case 0: SendSmallMessage(socket);
            break;
        case 1: SendLargeMessage(socket);
            break;
        case 2: SendMultipleMessages(socket);
            break;
        case 3: SendSplitMessage(socket);
            break;
        }
    }

    void SendSmallMessage(boost::asio::ip::tcp::socket& socket)
    {
        TransmitMessage(socket, RandomSize(m_config.small_message_min, m_config.small_message_max), "small");
    }

    void SendLargeMessage(boost::asio::ip::tcp::socket& socket)
    {
        TransmitMessage(socket, RandomSize(m_config.large_message_min, m_config.large_message_max), "large");
    }

    void SendMultipleMessages(boost::asio::ip::tcp::socket& socket)
    {
        for (int i = 0; i < m_config.multiple_message_count; ++i)
        {
            TransmitMessage(socket, RandomSize(m_config.multiple_message_min, m_config.multiple_message_max),
                        "multiple part " + std::to_string(i + 1));
        }
    }

    void SendSplitMessage(boost::asio::ip::tcp::socket& socket)
    {
        auto data = ConstructMessage(RandomSize(m_config.split_message_min, m_config.split_message_max));
        size_t split_point = data.size() / 2;

        boost::asio::write(socket, boost::asio::buffer(data.data(), split_point));
        std::cout << "Thread " << std::this_thread::get_id() << " sent split message part 1: " << split_point <<
            " bytes" << std::endl;

        std::this_thread::sleep_for(std::chrono::milliseconds(10));

        boost::asio::write(socket, boost::asio::buffer(data.data() + split_point, data.size() - split_point));
        std::cout << "Thread " << std::this_thread::get_id() << " sent split message part 2: " << (data.size() -
            split_point) << " bytes" << std::endl;
    }

    void TransmitMessage(boost::asio::ip::tcp::socket& socket, size_t payload_size, const std::string& message_type)
    {
        auto data = ConstructMessage(payload_size);
        boost::asio::write(socket, boost::asio::buffer(data));
        std::cout << "Thread " << std::this_thread::get_id() << " sent " << message_type << " message: " << data.size()
            << " bytes" << std::endl;
    }

    std::vector<std::byte> ConstructMessage(size_t payload_size)
    {
        Node node(payload_size, 0x01, 0x02);
        std::vector<std::byte> payload(payload_size);
        std::generate(payload.begin(), payload.end(), [&]() { return static_cast<std::byte>(m_gen()); });
        node.setPayload(payload.data(), payload_size);

        auto serialized = node.serialize();

        // 使用迭代器从指针和长度构造 vector
        std::vector<std::byte> result(serialized.buffer.get(), serialized.buffer.get() + serialized.length);
        return result;
    }


    void DelayBetweenMessages(std::default_random_engine& gen)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(
            std::uniform_int_distribution<>(m_config.delay_between_messages_min,
                                            m_config.delay_between_messages_max)(gen)));
    }

    size_t RandomSize(int min, int max)
    {
        return std::uniform_int_distribution<>(min, max)(m_gen);
    }

    std::string m_host;
    std::string m_port;
    Config m_config;
    std::mt19937 m_gen { std::random_device { }() };
};

int main()
{
    try
    {
        StressClient::Config config;
        config.thread_count = 4;
        config.messages_per_thread = 1000;
        config.total_test_duration = std::chrono::minutes(1);

        StressClient client("127.0.0.1", "10120", config);
        client.StartLoad();
    }
    catch (const std::exception& e)
    {
        std::cerr << "Exception: " << e.what() << std::endl;
    }

    return 0;
}
